# coding=utf-8
# -*- coding: UTF-8 -*-
#
# This file is part of the kernelstudio package.
#
# (c) 2014-2025 zlin <admin@kernelstudio.com>
#
# For the full copyright and license information, please view the LICENSE file
# that was distributed with this source code.
import logging
import os

from paddleocr import PaddleOCR

from ksai.document.ocr.config import ocr_executor_config
from ksai.document.ocr.executor import OCRExecutor
from ksai.utils.paths import get_project_models_path, get_project_storage_path, get_project_logs_path

logger = logging.getLogger(__name__)


class LocalOCRExecutor(OCRExecutor):

    def __init__(self, engine: PaddleOCR = None):
        super().__init__()
        self.engine = engine
        self._init_engine()

    def _init_engine(self):
        if self.engine is None:
            model_path = get_project_models_path()
            runtime_path = get_project_storage_path('_ocr')
            logs_path = get_project_logs_path()
            real_options = {
                'use_angle_cls': ocr_executor_config.OCR_USE_ANGLE_CLS,
                'lang': ocr_executor_config.OCR_LANG,
                'use_gpu': ocr_executor_config.OCR_USE_GPU,
                'det_model_dir': os.path.join(model_path, f"ocr/{ocr_executor_config.OCR_LANG}/det"),
                'rec_model_dir': os.path.join(model_path, f"ocr/{ocr_executor_config.OCR_LANG}/rec"),
                'cls_model_dir': os.path.join(model_path, f"ocr/{ocr_executor_config.OCR_LANG}/cls"),
                'crop_res_save_dir': os.path.join(runtime_path, 'crop'),
                'output': os.path.join(runtime_path, 'output'),
                'image_dir': os.path.join(runtime_path, 'image'),
                'save_log_path': os.path.join(logs_path)
            }

            logger.debug(f"OCR options: {real_options}")

            self.engine = PaddleOCR(**real_options)

    def execute(self, content: bytes) -> str:
        result = self.engine.ocr(content, cls=True)
        # 提取文本并合并
        text = " ".join([line[-1][0] for line in result[0]])
        return text
